﻿import numpy as np
import os
import torch
import torch.backends.cudnn as cudnn
import random

def set_seed(args):
    """设置随机种子以确保实验可重现"""
    print('Setting random seed for reproductivity..')
    
    # 处理不同类型的输入参数
    if isinstance(args, int):
        random_seed = args
    elif hasattr(args, 'random_seed'):
        random_seed = args.random_seed
    else:
        random_seed = 42  # 默认种子
    
    random.seed(random_seed)
    np.random.seed(random_seed)
    torch.manual_seed(random_seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(random_seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    os.environ['PYTHONHASHSEED'] = str(random_seed)
    print(f'[INFO] 随机种子已设置为: {random_seed}')
